import argparse
import config.config_subgroup as subgroup_config
import pandas as pd
import os
import numpy as np
from tqdm import tqdm   # 新增

parser = argparse.ArgumentParser()
parser.add_argument("--experiment", type=str, default="classification",
                    choices=["classification", "regression"])
parser.add_argument("--model", type=str, default="rf",
                    choices=["transformer","alex","res","mobile","mlp","lstm","rnn","rf","xgb","svm","linear","gbr"])
args = parser.parse_args()

# 读取亚组点估计并输出成完整的csv
print(f"Reading subgroup point estimates for {args.model} in {args.experiment}...")
subgroup_path=f"results_subgroup/{args.experiment}/{args.model}"
df=pd.DataFrame(columns=["acc","f1","precision","recall","roc_auc","cpa","infer_time","subgroup","n_samples", "variable"])
dataframes = []  # 列表收集所有DataFrame

# 外层：subgroup 进度条
for subgroup in tqdm(subgroup_config.SUBGROUPS, desc="Processing subgroups"):
    subgroup_dir=os.path.join(subgroup_path, subgroup)
    subgroup_df = [] # 一个亚组的所有文件
    
    # 内层：文件进度条
    for file in tqdm(os.listdir(subgroup_dir), desc=f"Files in {subgroup}", leave=False):
        if file.endswith(".csv"):
            print("Reading:", file)
            df_subgroup = pd.read_csv(os.path.join(subgroup_dir, file))
            print("Shape after read:", df_subgroup.shape)

            df_subgroup = df_subgroup.drop(columns=[subgroup], errors="ignore")
            print("After drop:", df_subgroup.shape)

            df_subgroup["variable"] = subgroup
            # df_subgroup = df_subgroup.iloc[1:]
            # print("After iloc:", df_subgroup.shape)  # <-- 看这里是不是 (0, n)

            metrics_to_eval = ["acc", "f1", "precision", "recall", "roc_auc"]

            # 先计算 SE
            for metric in metrics_to_eval:
                if metric in df_subgroup.columns:
                    p = df_subgroup[metric].astype(float)
                    n = df_subgroup["n_samples"].astype(float).replace(0, np.nan)

                    var = p * (1 - p) / n
                    se = np.sqrt(var.clip(lower=0))
                    df_subgroup[metric + "_se"] = se

            # 再统一调整列顺序
            cols = df_subgroup.columns.tolist()
            for metric in metrics_to_eval[::-1]:  # 倒序避免 index 问题
                se_col = metric + "_se"
                if se_col in cols:
                    metric_idx = cols.index(metric)
                    cols.insert(metric_idx + 1, cols.pop(cols.index(se_col)))
            df_subgroup = df_subgroup[cols]
            subgroup_df.append(df_subgroup)
    dataframes.append(pd.concat(subgroup_df, ignore_index=True))

df=pd.concat(dataframes, ignore_index=True)
df.to_csv(f"results_subgroup/{args.experiment}/{args.model}/subgroup_point_estimates.csv", index=False)

# 可视化成表格，包括变量列，，每个组占全样本的比例，组的样本数/全样本；acc, f1, precision, recall, roc_auc, 每个包括点估计和95%置信区间；将前面的数据画成森林图，数据+森林图穿插排列，每个指标都要